import cv2
import numpy as np
def read_image(path):
    img = cv2.imread(path, cv2.IMREAD_GRAYSCALE)
    if img is None:
        raise FileNotFoundError(f"图像 '{path}' 未找到")
    return img


def embed_watermark(carrier, watermark, bit_position):
    # 调整水印尺寸与载体一致
    watermark = cv2.resize(watermark, (carrier.shape[1], carrier.shape[0]))
    # 二值化水印
    watermark_bin = np.where(watermark > 100, 2 ** bit_position, 0).astype(np.uint8)
    # 预处理：清空相应位
    mask = np.full_like(carrier, 255 - 2 ** bit_position, dtype=np.uint8)
    carrier_cleared = cv2.bitwise_and(carrier, mask)
    # 嵌入水印
    return cv2.bitwise_or(carrier_cleared, watermark_bin)


def extract_watermark(img, bit_position):
    extracted = cv2.bitwise_and(img, np.full_like(img, 2 ** bit_position, dtype=np.uint8))
    return np.where(extracted > 0, 255, 0).astype(np.uint8)


def main():
    try:
        # 读取载体图像
        img = read_image("3/lena_gray.jpg")

        # 嵌入第一层水印
        watermark1 = read_image("4/treasure1.jpg")
        img_with_watermark1 = embed_watermark(img, watermark1, 0)

        # 嵌入第二层水印
        watermark2 = read_image("1.jpeg")
        img_with_both_watermarks = embed_watermark(img_with_watermark1, watermark2, 1)

        # 提取水印
        extracted1 = extract_watermark(img_with_both_watermarks, 0)
        extracted2 = extract_watermark(img_with_both_watermarks, 1)

        # 显示结果
        cv2.imshow("Original   Image", img)
        cv2.imshow("With   Both Watermarks", img_with_both_watermarks)
        cv2.imshow("Extracted   Watermark1 (LSB)", extracted1)
        cv2.imshow("Extracted   Watermark2 (Second LSB)", extracted2)

        # 保存结果
        cv2.imwrite("watermarked_lena.jpg", img_with_both_watermarks)
        cv2.imwrite("extracted_watermark1.jpg", extracted1)
        cv2.imwrite("extracted_watermark2.jpg", extracted2)

        cv2.waitKey(0)
        cv2.destroyAllWindows()

    except Exception as e:
        print(f"错误发生: {e}")
        cv2.destroyAllWindows()


if __name__ == "__main__":
    main()